Replace functools.partial(F.scaled_dot_product_attention, ...) with a
closure that resolves F.scaled_dot_product_attention at call time. This
ensures CP's runtime monkey-patch of the function is picked up by all
custom models instead of being bypassed by the early-bound reference.

Also make _attach_context_parallel_hooks public (renamed to
attach_context_parallel_hooks).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: hemildesai <hemild@nvidia.com>

…ends

Extract SDPA backend selection into a resolve_sdpa_method() helper that
accepts string names from YAML config (e.g. ["flash_attention",
"efficient_attention"]) and converts them to SDPBackend enum members.
When no explicit config is provided, auto-selects based on CP and
activation checkpointing constraints.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: hemildesai <hemild@nvidia.com>

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: hemildesai <hemild@nvidia.com>

Replace the assert that required all attention modules to be TE
DotProductAttention with a continue, so dense (SDPA) attention
modules are gracefully skipped. This allows MoE models to use
context parallelism with non-TE attention backends.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: hemildesai <hemild@nvidia.com>

Move the resolve_sdpa_method helper from train_ft.py to
_transformers/model_init.py per review feedback. The config
resolution (reading sdpa_method from YAML and passing it to
build_model) remains in train_ft.py.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: hemildesai <hemild@nvidia.com>

Move the attach_context_parallel_hooks call from train_ft.py into
apply_model_infrastructure in infrastructure.py, which already has
access to the device mesh. Add _uses_te_attention helper that inspects
the model's self_attn.attn_module instances to determine if TE
DotProductAttention is used, replacing the config-based check.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: hemildesai <hemild@nvidia.com>

When the model is an AutoPipeline, iterate over model.parts to inspect
self_attn modules instead of only the pipeline wrapper itself.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: hemildesai <hemild@nvidia.com>

Move the resolve_sdpa_method call from build_model in train_ft.py into
from_pretrained/from_config in auto_model.py, where device_mesh and
activation_checkpointing are available. train_ft.py now passes raw
string values from YAML as sdpa_method without pre-resolving.

resolve_sdpa_method now accepts both string names and SDPBackend enum
values, making the API flexible for both YAML config and programmatic
use.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: hemildesai <hemild@nvidia.com>